from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import json
import io
import numpy as np
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt

from libs.clean_structure import load_and_clean_json_data
from libs.extract_time import extract_time_domain_features
from libs.extract_frequency import extract_frequency_domain_features
from libs.stft_feature import stft_feature_extraction, wavelet_feature_extraction

app = Flask(__name__)
CORS(app)  # 添加CORS支持，允许所有域名访问

@app.route('/clean_data', methods=['POST'])
def clean_data():
    """
    数据清洗接口
    接收JSON数据，返回清洗后的数据和报告
    """
    try:
        data = request.get_json()
        if not data:
            return jsonify({"error": "未提供JSON数据"}), 400

        cleaned_data, report = load_and_clean_json_data(data)
        return jsonify({
            "cleaned_data": cleaned_data.to_dict(orient='records'),
            "report": report
        }), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_time_series', methods=['POST'])
def analyze_time_series():
    """
    时序分析接口
    接收时间序列数据，返回时域特征
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data:
            return jsonify({"error": "未提供信号数据"}), 400

        features = extract_time_domain_features(data["signal"])
        return jsonify({"features": features}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_time_series_plot', methods=['POST'])
def analyze_time_series_plot():
    """
    时序分析接口（图片版）
    接收时间序列数据，返回时域特征分析图
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data:
            return jsonify({"error": "未提供信号数据"}), 400

        signal = data["signal"]
        features = extract_time_domain_features(signal)
        
        # 创建图片
        plt.figure(figsize=(12, 8))
        
        # 绘制原始信号图
        plt.subplot(2, 2, 1)
        t = np.linspace(0, len(signal)/1000, len(signal))
        plt.plot(t, signal)
        plt.title('时间序列信号')
        plt.xlabel('时间 (s)')
        plt.ylabel('幅值')
        plt.grid(True)
        
        # 绘制统计特征线
        mean_val = features['mean_value']
        std_val = features['std_value']
        plt.axhline(y=mean_val, color='r', linestyle='--', label=f'均值: {mean_val:.2f}')
        plt.axhline(y=mean_val + std_val, color='g', linestyle=':', label=f'均值+标准差: {mean_val + std_val:.2f}')
        plt.axhline(y=mean_val - std_val, color='g', linestyle=':', label=f'均值-标准差: {mean_val - std_val:.2f}')
        plt.legend()
        
        # 绘制无量纲特征柱状图
        plt.subplot(2, 2, 2)
        dimensionless_features = {
            '峰值因子': features['crest_factor'],
            '脉冲因子': features['impulse_factor'],
            '波形因子': features['shape_factor'],
            '裕度因子': features['clearance_factor'],
            '偏度': features['skewness'],
            '峭度': features['kurtosis']
        }
        
        feature_names = list(dimensionless_features.keys())
        feature_values = list(dimensionless_features.values())
        
        bars = plt.bar(range(len(feature_names)), feature_values, color='skyblue')
        plt.xlabel('特征名称')
        plt.ylabel('特征值')
        plt.title('无量纲时域特征')
        plt.xticks(range(len(feature_names)), feature_names, rotation=45)
        
        # 在柱状图上添加数值标签
        for bar, value in zip(bars, feature_values):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height(), 
                     f'{value:.2f}', ha='center', va='bottom')
        
        # 绘制时域特征表格
        plt.subplot(2, 2, 3)
        plt.axis('off')
        feature_table_data = [
            ['特征', '数值'],
            ['均值', f"{features['mean_value']:.6f}"],
            ['标准差', f"{features['std_value']:.6f}"],
            ['均方根值', f"{features['rms_value']:.6f}"],
            ['最大值', f"{features['max_value']:.6f}"],
            ['峰峰值', f"{features['peak_to_peak']:.6f}"],
            ['绝对平均值', f"{features['abs_mean_value']:.6f}"]
        ]
        
        plt.table(cellText=feature_table_data[1:], colLabels=feature_table_data[0], 
                  cellLoc='center', loc='center', bbox=[0, 0, 1, 1])
        plt.title('时域特征值')
        
        # 绘制概率密度分布图
        plt.subplot(2, 2, 4)
        plt.hist(signal, bins=50, density=True, alpha=0.7, color='lightgreen')
        plt.xlabel('幅值')
        plt.ylabel('概率密度')
        plt.title('信号概率密度分布')
        plt.grid(True)
        
        plt.tight_layout()
        
        # 保存图片到内存
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format='png', dpi=150, bbox_inches='tight')
        img_buffer.seek(0)
        plt.close()
        
        return send_file(img_buffer, mimetype='image/png')
    except Exception as e:
        plt.close()
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_frequency', methods=['POST'])
def analyze_frequency():
    """
    频域分析接口
    接收时间序列数据，返回频域特征
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data or "sampling_rate" not in data:
            return jsonify({"error": "未提供信号数据或采样率"}), 400

        features, _ = extract_frequency_domain_features(data["signal"], data["sampling_rate"])
        return jsonify({"features": features}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_frequency_plot', methods=['POST'])
def analyze_frequency_plot():
    """
    频域分析接口（图片版）
    接收时间序列数据，返回频域特征分析图
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data or "sampling_rate" not in data:
            return jsonify({"error": "未提供信号数据或采样率"}), 400

        signal = data["signal"]
        sampling_rate = data["sampling_rate"]
        features, spectrum = extract_frequency_domain_features(signal, sampling_rate)
        
        # 创建图片
        plt.figure(figsize=(15, 10))
        
        # 绘制原始信号
        plt.subplot(2, 2, 1)
        t = np.linspace(0, len(signal)/sampling_rate, len(signal))
        plt.plot(t, signal)
        plt.title('时域信号')
        plt.xlabel('时间 (s)')
        plt.ylabel('振幅')
        plt.grid(True)
        
        # 绘制振幅频谱
        plt.subplot(2, 2, 2)
        plt.plot(spectrum['frequencies'], spectrum['amplitudes'])
        plt.title('振幅频谱')
        plt.xlabel('频率 (Hz)')
        plt.ylabel('振幅')
        plt.grid(True)
        
        # 标记主导频率
        dominant_freq = features['dominant_frequency']
        max_amp = features['max_amplitude']
        plt.plot(dominant_freq, max_amp, 'ro', label=f'主导频率: {dominant_freq:.1f}Hz')
        plt.legend()
        
        # 绘制功率谱密度 (PSD)
        plt.subplot(2, 2, 3)
        plt.semilogy(spectrum['frequencies_psd'], spectrum['psd'])
        plt.title('功率谱密度 (PSD)')
        plt.xlabel('频率 (Hz)')
        plt.ylabel('PSD (dB/Hz)')
        plt.grid(True)
        
        # 绘制特征值条形图
        plt.subplot(2, 2, 4)
        feature_names = list(features.keys())
        feature_values = list(features.values())
        plt.barh(range(len(feature_names)), feature_values)
        plt.yticks(range(len(feature_names)), feature_names)
        plt.title('频域特征值')
        plt.xlabel('特征值')
        plt.tight_layout()
        
        # 保存图片到内存
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format='png', dpi=150, bbox_inches='tight')
        img_buffer.seek(0)
        plt.close()
        
        return send_file(img_buffer, mimetype='image/png')
    except Exception as e:
        plt.close()
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_stft', methods=['POST'])
def analyze_stft():
    """
    STFT时频域分析接口
    接收时间序列数据，返回STFT特征
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data or "fs" not in data:
            return jsonify({"error": "未提供信号数据或采样频率"}), 400

        features, _ = stft_feature_extraction(data["signal"], data["fs"])
        # 只返回部分关键特征以减少数据量
        result = {
            "spectral_centroid": features["spectral_centroid"].tolist(),
            "spectral_bandwidth": features["spectral_bandwidth"].tolist(),
            "spectral_energy": features["spectral_energy"].tolist(),
            "time_points": features["time_points"].tolist()
        }
        return jsonify({"features": result}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_stft_plot', methods=['POST'])
def analyze_stft_plot():
    """
    STFT时频域分析接口（图片版）
    接收时间序列数据，返回STFT特征分析图
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data or "fs" not in data:
            return jsonify({"error": "未提供信号数据或采样频率"}), 400

        signal = data["signal"]
        fs = data["fs"]
        features, stft_data = stft_feature_extraction(signal, fs)
        
        # 创建图片
        plt.figure(figsize=(15, 10))
        
        # 绘制STFT谱图
        plt.subplot(2, 2, 1)
        plt.pcolormesh(stft_data['t'], stft_data['f'], np.abs(stft_data['Zxx']), shading='gouraud')
        plt.title('短时傅里叶变换(STFT)谱图')
        plt.xlabel('时间 [秒]')
        plt.ylabel('频率 [Hz]')
        plt.colorbar(label='幅度')
        
        # 绘制不同时刻的频谱切片
        plt.subplot(2, 2, 2)
        # 选择几个时间点的频谱切片
        time_indices = [0, len(stft_data['t'])//2, len(stft_data['t'])-1]
        labels = ['开始', '中间', '结束']
        colors = ['red', 'green', 'blue']
        
        for i, (idx, label, color) in enumerate(zip(time_indices, labels, colors)):
            spectrum_slice = np.abs(stft_data['Zxx'][:, idx])
            plt.plot(stft_data['f'], spectrum_slice, color=color, label=f'{label}时刻频谱')
        
        plt.title('不同时刻的频谱切片')
        plt.xlabel('频率 [Hz]')
        plt.ylabel('幅度')
        plt.legend()
        plt.grid(True)
        
        # 绘制谱质心变化图
        plt.subplot(2, 2, 3)
        plt.plot(features['time_points'], features['spectral_centroid'])
        plt.title('谱质心随时间变化')
        plt.xlabel('时间 [秒]')
        plt.ylabel('谱质心 [Hz]')
        plt.grid(True)
        
        # 绘制谱带宽变化图
        plt.subplot(2, 2, 4)
        plt.plot(features['time_points'], features['spectral_bandwidth'])
        plt.title('谱带宽随时间变化')
        plt.xlabel('时间 [秒]')
        plt.ylabel('谱带宽 [Hz]')
        plt.grid(True)
        
        plt.tight_layout()
        
        # 保存图片到内存
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format='png', dpi=150, bbox_inches='tight')
        img_buffer.seek(0)
        plt.close()
        
        return send_file(img_buffer, mimetype='image/png')
    except Exception as e:
        plt.close()
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_wavelet', methods=['POST'])
def analyze_wavelet():
    """
    小波变换时频域分析接口
    接收时间序列数据，返回小波变换特征
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data:
            return jsonify({"error": "未提供信号数据"}), 400

        features, _ = wavelet_feature_extraction(data["signal"])
        # 只返回部分关键特征以减少数据量
        result = {
            "energy_by_scale": features["energy_by_scale"].tolist(),
            "max_coeff_by_scale": features["max_coeff_by_scale"].tolist(),
            "mean_coeff_by_scale": features["mean_coeff_by_scale"].tolist(),
            "total_energy": float(features["total_energy"]),
            "scales": features["scales"].tolist()
        }
        return jsonify({"features": result}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/analyze_wavelet_plot', methods=['POST'])
def analyze_wavelet_plot():
    print("analyze_wavelet_plot算法调用")
    """
    小波变换时频域分析接口（图片版）
    接收时间序列数据，返回小波变换特征分析图
    """
    try:
        data = request.get_json()
        if not data or "signal" not in data:
            return jsonify({"error": "未提供信号数据"}), 400

        signal = data["signal"]
        features, wavelet_data = wavelet_feature_extraction(signal)
        
        # 创建图片
        plt.figure(figsize=(15, 10))
        
        # 绘制小波变换系数图
        plt.subplot(2, 2, 1)
        # 创建时间轴数据，与信号长度一致
        t = np.linspace(0, 1, np.abs(wavelet_data['coefficients']).shape[1])
        plt.pcolormesh(t, wavelet_data['frequencies'], 
                       np.abs(wavelet_data['coefficients']), shading='gouraud')
        plt.title('小波变换系数')
        plt.xlabel('时间')
        plt.ylabel('频率')
        plt.colorbar(label='幅度')
        
        # 绘制能量随尺度变化图
        plt.subplot(2, 2, 2)
        plt.plot(features['scales'], features['energy_by_scale'])
        plt.title('能量随尺度的变化')
        plt.xlabel('尺度')
        plt.ylabel('能量')
        plt.grid(True)
        
        # 绘制各尺度最大系数图
        plt.subplot(2, 2, 3)
        plt.plot(features['scales'], features['max_coeff_by_scale'])
        plt.title('各尺度最大系数')
        plt.xlabel('尺度')
        plt.ylabel('最大系数')
        plt.grid(True)
        
        # 绘制小波系数幅度分布直方图
        plt.subplot(2, 2, 4)
        plt.hist(np.abs(wavelet_data['coefficients']).flatten(), bins=50, alpha=0.7)
        plt.title('小波系数幅度分布')
        plt.xlabel('幅度')
        plt.ylabel('频次')
        plt.grid(True)
        
        plt.tight_layout()
        
        # 保存图片到内存
        img_buffer = io.BytesIO()
        plt.savefig(img_buffer, format='png', dpi=150, bbox_inches='tight')
        img_buffer.seek(0)
        plt.close()
        
        return send_file(img_buffer, mimetype='image/png')
    except Exception as e:
        plt.close()
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    from dotenv import load_dotenv
    import os
    load_dotenv()
    app.run(debug=True, host=os.getenv('Host', '0.0.0.0'), port=int(os.getenv('Port', 5000)))